import random
import numpy as np

class GeneticAlgorithm8Queens:
    def __init__(self, population_size=100, mutation_rate=0.1, max_generations=1000):
        self.population_size = population_size
        self.mutation_rate = mutation_rate
        self.max_generations = max_generations
        self.board_size = 8
        
    def create_individual(self):
        """
        创建一个个体（解），表示棋盘上皇后的位置
        每个索引代表行，值代表列
        """
        return list(np.random.permutation(self.board_size))
    
    def create_population(self):
        """
        创建初始种群
        """
        return [self.create_individual() for _ in range(self.population_size)]
    
    def fitness(self, individual):
        """
        计算适应度函数 - 互相攻击的皇后对数越少，适应度越高
        """
        conflicts = 0
        for i in range(self.board_size):
            for j in range(i+1, self.board_size):
                # 检查是否在同一列
                if individual[i] == individual[j]:
                    conflicts += 1
                # 检查是否在对角线上
                if abs(individual[i] - individual[j]) == abs(i - j):
                    conflicts += 1
        return 1 / (1 + conflicts)
    
    def selection(self, population):
        """
        选择操作 - 使用轮盘赌选择
        """
        fitness_scores = [self.fitness(individual) for individual in population]
        total_fitness = sum(fitness_scores)
        
        if total_fitness == 0:
            # 如果所有个体适应度都为0，则随机选择
            return random.choice(population)
            
        probabilities = [score/total_fitness for score in fitness_scores]
        selected_index = np.random.choice(len(population), p=probabilities)
        return population[selected_index]
    
    def crossover(self, parent1, parent2):
        """
        交叉操作 - 使用顺序交叉(OX1)
        """
        start = random.randint(0, self.board_size-2)
        end = random.randint(start+1, self.board_size-1)
        
        child1 = [-1] * self.board_size
        child2 = [-1] * self.board_size
        
        # 复制片段
        child1[start:end+1] = parent1[start:end+1]
        child2[start:end+1] = parent2[start:end+1]
        
        # 填充剩余位置
        self.fill_child(child1, parent2, start, end)
        self.fill_child(child2, parent1, start, end)
        
        return child1, child2
    
    def fill_child(self, child, parent, start, end):
        """
        辅助函数：填充子代中未确定的位置
        """
        pointer = (end + 1) % self.board_size
        for i in range(end+1, self.board_size):
            if parent[i] not in child:
                child[pointer] = parent[i]
                pointer = (pointer + 1) % self.board_size
                
        for i in range(0, end+1):
            if parent[i] not in child:
                child[pointer] = parent[i]
                pointer = (pointer + 1) % self.board_size
    
    def mutate(self, individual):
        """
        变异操作 - 随机交换两个位置
        """
        if random.random() < self.mutation_rate:
            i, j = random.sample(range(self.board_size), 2)
            individual[i], individual[j] = individual[j], individual[i]
        return individual
    
    def solve(self):
        """
        主算法流程
        """
        # 创建初始种群
        population = self.create_population()
        
        for generation in range(self.max_generations):
            # 检查是否有解
            for individual in population:
                if self.fitness(individual) == 1.0:  # 无冲突，找到了解
                    return individual, generation
            
            # 创建新一代
            new_population = []
            
            # 保留最佳个体（精英主义）
            fitness_scores = [(self.fitness(ind), ind) for ind in population]
            fitness_scores.sort(reverse=True)
            new_population.append(fitness_scores[0][1])  # 添加最佳个体
            
            # 生成其余个体
            while len(new_population) < self.population_size:
                parent1 = self.selection(population)
                parent2 = self.selection(population)
                
                child1, child2 = self.crossover(parent1, parent2)
                
                child1 = self.mutate(child1)
                child2 = self.mutate(child2)
                
                new_population.extend([child1, child2])
            
            # 更新种群
            population = new_population[:self.population_size]
        
        # 返回最佳解
        best_individual = max(population, key=self.fitness)
        return best_individual, self.max_generations
    
    def print_board(self, individual):
        """
        打印棋盘
        """
        print("\n+---+---+---+---+---+---+---+---+")
        for row in range(self.board_size):
            print("|", end="")
            for col in range(self.board_size):
                if individual[row] == col:
                    print(" Q ", end="|")
                else:
                    print("   ", end="|")
            print("\n+---+---+---+---+---+---+---+---+")

# 主程序
if __name__ == "__main__":
    ga = GeneticAlgorithm8Queens(population_size=100, mutation_rate=0.1, max_generations=1000)
    
    print("正在使用遗传算法解决8皇后问题...")
    solution, generations = ga.solve()
    
    print(f"\n找到解，共用了 {generations} 代")
    print(f"解的适应度: {ga.fitness(solution)}")
    print("皇后位置 (行, 列):")
    for i, col in enumerate(solution):
        print(f"  第{i+1}行: 第{col+1}列")
    
    ga.print_board(solution)